import argparse
import io
from time import sleep
import grpc
import numpy as np
from PIL import Image
# import imageio.v3 as iio  # imageio v3 API
import imageio

from google.protobuf import empty_pb2

# 如果你把 stubs 生成在包里：
import play_with_pybullet.grpc_env.arm_env_pb2  as pb
import play_with_pybullet.grpc_env.arm_env_pb2_grpc as rpc
# 如果是“扁平导入”，改成：
# import arm_env_pb2 as pb
# import arm_env_pb2_grpc as rpc


def main(args):
    channel = grpc.insecure_channel(f"{args.host}:{args.port}")
    stub = rpc.ArmEnvStub(channel)

    # 1) 查询空间
    spaces = stub.GetSpaces(empty_pb2.Empty())
    print("Action dim:", spaces.action_dim,
          "low:", list(spaces.action_low),
          "high:", list(spaces.action_high))
    print("Obs dim:", spaces.obs_dim)

    # 2) reset（决定是否启用渲染）
    reset = pb.ResetRequest(
        render_mode="rgb_array" if args.render else "none",
        has_seed=(args.seed is not None),
        seed=(args.seed or 0)
    )
    obs = stub.Reset(reset).observation
    print("Initial obs:", obs)

    # 3) 初始化视频写入器（固定分辨率 & 帧率）
    writer = None
    if args.save_video:
        # 注意：所有帧尺寸必须一致，使用请求宽高进行强制 resize
        writer = imageio.get_writer(
            args.video_out,
            fps=args.fps,
            codec=args.codec,          # "libx264" 或 "mpeg4"
            format="ffmpeg",
            ffmpeg_params=["-pix_fmt", "yuv420p", "-movflags", "faststart"]
        )
        print(f"[Video] writing to {args.video_out} @ {args.fps} fps, codec={args.codec}")

    # 4) rollout
    rng = np.random.default_rng(0)
    try:
        for t in range(args.steps):
            a = rng.uniform(-1.0, 1.0, size=(spaces.action_dim,)).astype(np.float32)
            step = stub.Step(pb.StepRequest(action=list(a)))
            print(f"t={t:03d} reward={step.reward:.4f} term={step.terminated} "
                  f"trunc={step.truncated} obs0={step.observation[:4]}")

            # 采帧：每 render_every 步取一帧
            if args.render and (t % args.render_every == 0):
                frame = stub.Render(pb.RenderRequest(width=args.width, height=args.height))
                # PNG bytes -> RGB ndarray(H,W,3)
                pil = Image.open(io.BytesIO(frame.rgb_png)).convert("RGB")
                if args.save_frames:  # 可选同时落单帧
                    pil.save(f"{args.frames_dir}/frame_{t:04d}.png")
                if writer is not None:
                    writer.append_data(np.asarray(pil, dtype=np.uint8))

            # 可选：按 “仿真节拍” 进行节流（与视频 fps 无强制关系）
            if args.rt_sleep > 0:
                sleep(args.rt_sleep)
    finally:
        if writer is not None:
            writer.close()
        # stub.Close(empty_pb2.Empty())  # 如果服务端实现了 Close 再启用
        channel.close()


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=50051)
    p.add_argument("--steps", type=int, default=500)
    p.add_argument("--seed", type=int, default=None)

    # 渲染/取帧
    p.add_argument("--render", action="store_true", default=True)
    p.add_argument("--render-every", type=int, default=1)
    p.add_argument("--width", type=int, default=640)
    p.add_argument("--height", type=int, default=480)

    # 视频参数
    p.add_argument("--save-video", action="store_true", default=True)
    p.add_argument("--video-out", type=str, default="out.mp4")
    p.add_argument("--fps", type=float, default=30.0)
    p.add_argument("--codec", type=str, default="libx264")  # 若环境不含 x264，可改 "mpeg4"

    # 可选：也同时落单帧
    p.add_argument("--save-frames", action="store_false", default=False)
    p.add_argument("--frames-dir", type=str, default="./frames")

    # 仿真节拍（控制 Step 速度；和最终视频 fps 独立）
    p.add_argument("--rt-sleep", type=float, default=0.0)

    main(p.parse_args())
